import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# Read the Excel file
file_path = 'Phi_Psi_angles.xlsx'  # Update this path with the actual path to your Excel file

# List of sheet names
sheet_names = ['Chain A', 'Chain B', 'Chain C', 'Chain D']  # Update these with the actual sheet names

# Extracting columns for each measurement
measurements = {
    'R375K': (1, 2),    # Column indices (0-based) for phi and psi angles
    'RW_LF': (3, 4),
    'W376A': (5, 6),
    'W376F': (7, 8)
}

# Define markers for each dataset
markers = {
    'R375K': 'o',  # Circle
    'RW_LF': 's',  # Square
    'W376A': '^',  # Triangle
    'W376F': 'D'   # Diamond
}

# Ensure the column names in `measurements` match the printed column names
colors = ['k', 'b', 'g', 'm']

for sheet_name in sheet_names:
    df = pd.read_excel(file_path, sheet_name=sheet_name)
    
    # Strip leading/trailing spaces from column names
    df.columns = df.columns.str.strip()

    # Print column names for debugging
    print(f"Column names in the DataFrame for {sheet_name}:")
    print(df.columns)
    
    plt.figure(figsize=(8, 8))

    for (label, (phi_idx, psi_idx)), color in zip(measurements.items(), colors):
        if phi_idx >= len(df.columns) or psi_idx >= len(df.columns):
            print(f"Skipping {label} for {sheet_name} because the columns are missing.")
            continue
        
        # Convert to numeric, coercing errors to NaN
        phi_angles = pd.to_numeric(df.iloc[:, phi_idx], errors='coerce')
        psi_angles = pd.to_numeric(df.iloc[:, psi_idx], errors='coerce')
        
        # Drop NaN values
        valid_indices = phi_angles.notna() & psi_angles.notna()
        phi_angles = phi_angles[valid_indices]
        psi_angles = psi_angles[valid_indices]
        
        # Customize data points with different shapes and black edges
        plt.scatter(phi_angles, psi_angles, s=50, color=color, marker=markers[label], edgecolor='black', linewidths=0.5, alpha=0.7, label=label)

    plt.xlim(-180, 180)
    plt.ylim(-180, 180)
    plt.xlabel('Phi (φ)', fontsize=14)  # Adjust font size here
    plt.ylabel('Psi (ψ)', fontsize=14)  # Adjust font size here
    plt.title(f'Ramachandran Plot - {sheet_name}', fontsize=16)  # Adjust font size here

    # Increase tick label font size for both axes
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)

    # Add gridlines crossing at 0 degrees from both axes
    plt.axhline(0, color='gray', linewidth=0.5)
    plt.axvline(0, color='gray', linewidth=0.5)

    # Place the legend outside the plot
    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5), fontsize=12)  # Adjust font size here
    plt.savefig(f'ramachandran_plot_{sheet_name}.png', bbox_inches='tight', dpi=600)
    plt.show()
